import gymnasium as gym
from gymnasium import spaces
import numpy as np
from decimal import Decimal
from backtest.backtest import BackTest
from backtest.order_book_sh import OrderBookSH
from backtest.stock_order import StockOrder
from typing import Dict, List

class TradingEnv(gym.Env):
    """
    股票拆单算法训练环境
    目标：将大单拆分成小单执行，平衡交易完成度和交易成本
    """
    
    def __init__(self, csv_path: str, symbol: str, order_log_file: str = "orders.log", 
                 target_quantity: int = 100000, max_steps: int = 6000, 
                 direction: str = "buy", initial_cash: float = 100000000.0, price_log_file: str = "logs/train/price.csv", is_train: bool = True):
        super().__init__()
        
        # 初始化核心组件
        self.order_book = OrderBookSH()
        self.order_book.set_symbol(symbol)
        self.backtest = BackTest(self.order_book, csv_path, log_path=price_log_file)
        
        # 环境参数
        self.current_step = 0
        self.max_steps = max_steps
        self.symbol = symbol
        self.direction = direction  # "buy" 或 "sell"
        self.initial_cash = initial_cash
        self.current_cash = initial_cash
        
        # 拆单目标参数
        self.target_quantity = target_quantity
        self.initial_target = target_quantity
        self.executed_quantity = 0
        self.executed_value = 0.0
        self.start_price = self.order_book.get_price()
        self.remaining_quantity = target_quantity
        
        # 观察空间：47维市场数据 + 8维拆单状态 = 55维
        self.observation_space = spaces.Box(
            low=0, high=1e11, shape=(55,), dtype=np.float32
        )
        
        # 动作空间：33个离散动作 (3种操作类型 × 11个数量等级)
        self.action_space = spaces.Discrete(33)
        
        # 状态变量
        self.__orderid_counter = -1
        # 当前环境是否是训练模式
        self.is_train = is_train
        
        # 奖励权重配置
        self.reward_weights = {
            'completion': 2.0,      # 完成度奖励
            'cost_penalty': -2.0,   # 成本惩罚
            'completion_bonus': 5.0, # 完成奖励
            'time_penalty': -1    # 时间惩罚
        }
        
        # 日志文件
        self.order_log_file = order_log_file
        with open(self.order_log_file, "w") as f:
            f.write(f"Target: {target_quantity} shares, Direction: {direction}, Initial Cash: {initial_cash}\n")

    def reset(self, *, seed=None, options=None):
        """重置环境状态"""
        super().reset(seed=seed)
        
        # 重置核心组件
        self.backtest.reset()
        self.current_step = 0
        self.backtest.step_once()
        self.__orderid_counter = -1
        
        # 重置拆单状态
        self.executed_quantity = 0
        self.executed_value = 0.0
        self.current_cash = self.initial_cash
        self.remaining_quantity = self.target_quantity

        
        self.start_price = self.order_book.get_price()
        
        # 获取初始观察值
        obs = self._get_observation()
        
        # 记录重置日志
        # self._log_event("RESET", f"Start price: {self.start_price}, Target: {self.target_quantity}")
        
        return obs, {}

    def step(self, action: int):
        pre_executed = self.executed_quantity
        pre_value = self.executed_value
        pre_remaining = self.remaining_quantity
    
        # 执行动作
        current_orders = self._execute_action(action)
    
        # 推进行情
        self.backtest.step_once()
        self.current_step += 1
    
        # 更新执行状态
        self._update_execution_status(current_orders)
    
        # 获取观察值
        obs = self._get_observation()
    
        # 计算奖励
        reward = self._calculate_comprehensive_reward(
            pre_executed, pre_value, pre_remaining, current_orders
        )
    
        # 检查终止
        terminated, truncated, info = self._check_termination()
    
        # === 添加额外指标到 info ===
        info.update({
            'completion_rate': self.executed_quantity / self.initial_target,
            'avg_price': self.executed_value / self.executed_quantity if self.executed_quantity > 0 else 0,
            'remaining_quantity': self.remaining_quantity,
            'executed_value': self.executed_value,
            'current_cash': self.current_cash,
            'market_price': float(self.order_book.get_price()),
            'vwap': self.executed_value / self.executed_quantity if self.executed_quantity > 0 else 0
        })
    
        # 记录日志
        # self._log_step_info(action, reward, current_orders, info)
    
        return obs, reward, terminated, truncated, info



    def _execute_action(self, action: int) -> List[StockOrder]:
        """执行交易动作"""
        action_type = action // 11  # 0=不动, 1=买, 2=卖
        amount_level = action % 11  # 0~10 表示数量等级
        
        current_orders = []
        
        # 检查是否应该执行动作
        if action_type == 0 or amount_level == 0 or self.remaining_quantity <= 0:
            return current_orders
        
        # 计算订单数量（不超过剩余数量）
        base_quantity = amount_level * 100
        quantity = min(base_quantity, self.remaining_quantity)
        
        # 生成订单
        order_id = self.__orderid_counter
        self.__orderid_counter -= 1
        
        # 确定价格和方向
        if self.direction == "buy":
            price = self.order_book.get_sellone_price()  # 买方向卖一价下单
            is_buy = True
        else:
            price = self.order_book.get_buyone_price()   # 卖方向买一价下单
            is_buy = False
        
        # 检查资金是否足够（对于买单）
        if is_buy and self.current_cash < quantity * price:
            quantity = int(self.current_cash / price)  # 调整数量以适应资金
            if quantity <= 0:
                return current_orders
        
        create_time = self.order_book.get_timestamp()
        
        # 创建订单
        order = StockOrder(order_id, self.symbol, quantity, price, is_buy, create_time)
        self.backtest.place_order(order)
        self._log_order(order)
        current_orders.append(order)
        
        return current_orders

    def _update_execution_status(self, orders: List[StockOrder]):
        """更新订单执行状态"""
        for order in orders:
            # 这里需要根据您的backtest实现获取实际执行情况
            # 暂时使用模拟执行结果
            if order.is_buy_order():
                # 买单执行逻辑
                executed_qty = min(order.get_filled_quantity(), self.remaining_quantity)
                executed_price = float(order.get_price())
                cost = executed_qty * executed_price
                
                if executed_qty > 0 and cost <= self.current_cash:
                    self.executed_quantity += executed_qty
                    self.executed_value += cost
                    self.current_cash -= cost
                    self.remaining_quantity -= executed_qty
            else:
                # 卖单执行逻辑
                executed_qty = order.get_filled_quantity()
                executed_price = order.get_price()
                revenue = executed_qty * executed_price
                
                self.executed_quantity += executed_qty
                self.executed_value += revenue
                self.current_cash += revenue
                self.remaining_quantity -= executed_qty

    def _calculate_comprehensive_reward(self, pre_executed, pre_value, pre_remaining, orders):
        rewards = {}

        # 1. 分步完成奖励（每完成 10% 一次性奖励）
        delta_completion = self.executed_quantity - pre_executed
        step_bonus = (self.executed_quantity // (0.1 * self.initial_target) -
                      pre_executed // (0.1 * self.initial_target))
        rewards['step_completion'] = float(step_bonus)

        # 2. 成本奖励：成交价格接近市场均价越好奖励越高
        if self.executed_quantity > 0:
            avg_price = self.executed_value / self.executed_quantity
            market_price = float(self.order_book.get_price())  # 转成 float
            price_diff_ratio = abs(avg_price - market_price) / market_price
            rewards['cost_reward'] = -np.tanh(5 * price_diff_ratio)
        else:
            rewards['cost_reward'] = 0.0

        # 3. 时间惩罚：指数形式更敏感
        time_ratio = self.current_step / self.max_steps
        rewards['time_penalty'] = -np.exp(time_ratio) + 1

        # 4. 剩余目标惩罚：未完成越多惩罚越大
        remaining_ratio = self.remaining_quantity / self.initial_target
        rewards['remaining_penalty'] = -remaining_ratio ** 2

        # 5. 完成目标 bonus
        rewards['completion_bonus'] = 5.0 if self.executed_quantity >= self.initial_target else 0.0

        # 权重配置
        weights = {
            'step_completion': 2.0,
            'cost_reward': 2.0,
            'time_penalty': 1.0,
            'remaining_penalty': 1.5,
            'completion_bonus': 3.0
        }

        total_reward = sum(rewards[key] * weights.get(key, 1.0) for key in rewards)
        return float(np.tanh(total_reward))

    
    def _check_termination(self):
        """检查终止条件并返回 info"""
        if self.is_train:
            # ✅ 训练模式下的终止条件
            terminated = bool(self.executed_quantity >= self.initial_target)
            truncated = bool(
                self.current_step >= self.max_steps or 
                self.backtest.is_done() or 
                self.remaining_quantity <= 0
            )
        else:
            # ✅ 预测/评估模式下的终止条件
            terminated = False
            truncated = bool(self.backtest.is_done())

        # 记录每步 info
        info = {
            'executed_quantity': self.executed_quantity,
            'completion_rate': self.executed_quantity / self.initial_target,
            'avg_price': self.executed_value / self.executed_quantity if self.executed_quantity > 0 else 0,
            'remaining_quantity': self.remaining_quantity,
            'current_cash': self.current_cash,
            'current_step': self.current_step,
            'vwap': self.executed_value / self.executed_quantity if self.executed_quantity > 0 else 0,
            'market_price': float(self.order_book.get_price())
        }

        return terminated, truncated, info

    def _get_observation(self):
        """获取环境观察值 - 修复列表转数组问题"""
        # 基础市场数据
        snapshot = self.order_book.get_snapshot()
        base_obs = snapshot.to_numpy()

        # 如果base_obs是列表，先转换为numpy数组
        if isinstance(base_obs, list):
            base_obs = np.array(base_obs, dtype=np.float32)
        else:
            # 确保是float32类型
            base_obs = base_obs.astype(np.float32)

        # 如果基础数据不是47维，需要调整
        if base_obs.shape[0] != 47:
            # 调整到47维
            if base_obs.shape[0] > 47:
                base_obs = base_obs[:47]
            else:
                base_obs = np.pad(base_obs, (0, 47 - base_obs.shape[0]), mode='constant')

        # 拆单状态特征
        completion_ratio = np.float32(self.executed_quantity / self.initial_target)
        remaining_ratio = np.float32(self.remaining_quantity / self.initial_target)
        time_ratio = np.float32(self.current_step / self.max_steps)

        current_price = np.float32(self.order_book.get_price())
        start_price_val = np.float32(self.start_price) if self.start_price else np.float32(1.0)
        price_ratio = np.float32(current_price / start_price_val)

        # 资金状态
        cash_ratio = np.float32(self.current_cash / self.initial_cash)

        # 组合观察值
        additional_features = np.array([
            completion_ratio,
            remaining_ratio, 
            time_ratio,
            price_ratio,
            cash_ratio,
            np.float32(self.executed_quantity),
            np.float32(self.remaining_quantity),
            np.float32(self.current_step)
        ], dtype=np.float32)

        # 合并观察值
        obs = np.concatenate([base_obs, additional_features])

        # 最终确保是float32
        return obs.astype(np.float32)

    def _log_order(self, order: StockOrder):
        """记录订单日志"""
        with open(self.order_log_file, "a") as f:
            direction = "buy" if order.is_buy_order() else "sell"
            f.write(f"step {self.current_step}: timestamp {self.order_book.get_timestamp()}: {direction} {order.get_quantity()}: price {order.get_price():.2f}\n")

    def _log_event(self, event_type: str, message: str):
        """记录事件日志"""
        with open(self.order_log_file, "a") as f:
            f.write(f"[{event_type}] {message}\n")

    def _log_step_info(self, action: int, reward: float, orders: List[StockOrder], info: Dict):
        """记录步骤信息"""
        with open(self.order_log_file, "a") as f:
            f.write(f"--- Step {self.current_step} Summary ---\n")
            f.write(f"Action: {action}, Reward: {reward:.4f}\n")
            f.write(f"Executed: {info['executed_quantity']}/{self.initial_target} ")
            f.write(f"({info['completion_rate']*100:.1f}%)\n")
            f.write(f"Avg Price: {info['avg_price']:.2f}, Market Price: {info['market_price']:.2f}\n")
            f.write(f"Remaining: {info['remaining_quantity']}, Cash: {info['current_cash']:.2f}\n")
            f.write(f"Orders this step: {len(orders)}\n\n")

    def render(self, mode="human"):
        """渲染环境状态"""
        if mode == "human":
            print(f"\n=== Step {self.current_step} ===")
            print(f"Target: {self.initial_target}, Executed: {self.executed_quantity}")
            print(f"Completion: {self.executed_quantity/self.initial_target*100:.1f}%")
            print(f"Remaining: {self.remaining_quantity}")
            if self.executed_quantity > 0:
                print(f"Avg Price: {self.executed_value/self.executed_quantity:.2f}")
            print(f"Market Price: {self.order_book.get_price():.2f}")
            print(f"Cash: {self.current_cash:.2f}")
            print(f"Progress: {self.current_step}/{self.max_steps} steps")

    def close(self):
        
        # 记录结束日志
        self._log_event("CLOSE", f"Final executed: {self.executed_quantity}/{self.initial_target}")